#include "LDPC.h"
#include <chrono> //timing
#include <cstdio> // sprintf
#include <iostream>
#include <random>

#define print_x(x) print_vec(x, #x, __func__, __LINE__)
void print_vec(const std::vector<bool> &v, const char *name, const char *func,
               int line) {
    fprintf(stdout, "[%s:%d] %s(%lu) ", func, line, name, v.size());
    for (auto e : v)
        fprintf(stdout, "%d ", e);
    fprintf(stdout, "\n");
}
void print_vec(const std::vector<double> &v, const char *name, const char *func,
               int line) {
    fprintf(stdout, "[%s:%d] %s(%lu) ", func, line, name, v.size());
    for (auto e : v)
        fprintf(stdout, "%f ", e);
    fprintf(stdout, "\n");
}

std::vector<double> lin_space(double start, double end, int num) {
    // catch rarely, throw often
    assert(num >= 2 && "The third parameter must be a positive integer >= 2!");

    int partitions = num - 1;
    std::vector<double> pts;
    // length of each segment
    double length = (end - start) / partitions;
    // first, not to change
    pts.push_back(start);
    for (int i = 1; i < num - 1; i++) {
        pts.push_back(start + i * length);
    }
    // last, not to change
    pts.push_back(end);
    return pts;
}

void runLDPC() {
    uint32_t nMaxIter = 8; // number of decoders
    int N = 1000;          // code word length
    double rate = 1.0 / 3; // code rate

    uint32_t K = uint32_t(std::ceil(N * rate)); // information length

    // instantiates a POLAR object
    nrLDPC ldpc = nrLDPC(K, rate);

    // random engines
    std::default_random_engine random_engine;
    std::bernoulli_distribution bern_dist;
    std::normal_distribution<double> norm_dist(0, 1);

    // Running parameters
    std::vector<double> EsN0_dB = lin_space(-4, -2, 9);
    std::vector<double> N0(EsN0_dB.size(), 0);
    print_x(EsN0_dB);
    print_x(N0);
    std::transform(EsN0_dB.begin(), EsN0_dB.end(), N0.begin(),
                   [](const float &x) { return std::pow(10.0, -x / 10.0); });
    print_x(EsN0_dB);
    print_x(N0);

    std::vector<double> ber(N0.size(), 0), bler(N0.size(), 0);
    std::vector<uint32_t> n_bit_errs(N0.size(), 0), n_blk_errs(N0.size(), 0);

    uint32_t n_max_blks = 10000;
    std::vector<bool> fillers(ldpc.getFillerLength(), 0);
    // loop each SNR
    for (uint32_t i = 0; i < N0.size(); i++) {
        // print progress
        fprintf(stdout, "\nNow running EsN0: %.2f dB [%d of %lu]", EsN0_dB[i],
                i + 1, N0.size());
        uint32_t print_len = 0;

        uint32_t n_blks_done = 0;
        clock_t tStart = clock(); // timing

        while ((n_blks_done < n_max_blks) && (n_blk_errs[i] < 100)) {
            // generate random bit stream
            // = {
            // 1,0,0,1,1,0,1,0,0,0,0,0,1,0,1,1,1,0,0,0,0,1,1,0,1,0,1,1,1,1,0,1,0,1
            // };
            std::vector<bool> msg;
            msg.reserve(K);
            for (uint32_t j = 0; j < K; j++)
                msg.push_back(bern_dist(random_engine));
            // print_x(msg);

            // add filler bits
            std::vector<bool> extMsg = msg;
            extMsg.insert(extMsg.end(), fillers.begin(), fillers.end());
            // print_x(extMsg);

            // LDPC encoding
            std::vector<bool> enc = ldpc.encode(extMsg);
            assert(ldpc.checkSumCodeWord(enc));
            // print_x(enc);

            // rate matching
            std::vector<bool> rm_enc = ldpc.rateMatch(enc, N);
            // print_x(rm_enc);

            // BPSK + AWGN
            std::vector<double> r;
            r.reserve(N);
            for (auto e : rm_enc)
                r.push_back(1 - 2.0 * e +
                            sqrt(N0[i] / 2.0) * norm_dist(random_engine));

            // compute soft bits as LLR
            std::vector<double> llr;
            llr.reserve(N);
            for (auto e : r)
                llr.push_back(4.0 * e / N0[i]);
            // print_x(llr);

            // rate recovery
            std::vector<double> rr_llr = ldpc.rateRecover(llr);
            // print_x(rr_llr);

            // scl decoding
            std::vector<bool> msg_cap = ldpc.decode(rr_llr, nMaxIter);
            // print_x(msg_cap);

            // count errors
            uint32_t n_errs = 0;
            for (uint32_t j = 0; j < K; j++) {
                if (msg[j] != msg_cap[j])
                    n_errs++;
            }

            if (n_errs) {
                n_bit_errs[i] += n_errs;
                n_blk_errs[i]++;
            }

            n_blks_done += 1;

            ber[i] = n_bit_errs[i] * 1.0 / K / n_blks_done;
            bler[i] = n_blk_errs[i] * 1.0 / n_blks_done;

            // print progress for every 10 blocks
            if (n_blks_done % 10 == 0 || n_blks_done == 1) {
                fprintf(stdout,
                        "Elapsed time: %.1f seconds, # tx blocks: %d,# "
                        "error blocks:%d, ber: %.5f, bler %.5f\n",
                        double(clock() - tStart) / CLOCKS_PER_SEC, n_blks_done,
                        n_blk_errs[i], ber[i], bler[i]);
            }
        }

        // print  progress when one SNR is finished
        fprintf(
            stdout,
            "Elapsed time: %.1f seconds, # tx blocks: %d,# error blocks:%d, "
            "ber: %.5f, bler %.5f\n",
            static_cast<double>(clock() - tStart) / CLOCKS_PER_SEC, n_blks_done,
            n_blk_errs[i], ber[i], bler[i]);
    }

    // print simulation result
    std::cout << std::endl;
    std::cout << "Modulation:"
              << "BPSK" << std::endl;
    std::cout << "[N,R] = [ " << N << "," << rate << "]" << std::endl;
    std::cout << "EsN0_dB = [";
    for (auto e : EsN0_dB)
        std::cout << e << " ";
    std::cout << "]" << std::endl;

    std::cout << "BER = [";
    for (auto e : ber)
        std::cout << e << " ";
    std::cout << "]" << std::endl;

    std::cout << "BLER = [";
    for (auto e : bler)
        std::cout << e << " ";
    std::cout << "]" << std::endl;
}

int main() {
    std::cout << "Hello, 5G-NR!" << std::endl;

    runLDPC();

    return 0;
}